package com.service.impl;

import com.alibaba.fastjson.JSONObject;
import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.model.auto.AiQuestion;
import com.model.auto.AiQuestionItem;
import com.model.auto.vo.*;
import com.service.IAiQuestionItemService;
import com.service.IAiQuestionService;
import com.service.IBaiduAiService;
import lombok.extern.slf4j.Slf4j;
import okhttp3.*;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.http.HttpHeaders;
import org.springframework.stereotype.Service;
import org.springframework.web.reactive.function.BodyInserters;
import org.springframework.web.reactive.function.client.WebClient;
import org.springframework.web.reactive.function.client.WebClientResponseException;
import reactor.core.publisher.Flux;

import java.io.IOException;
import java.time.LocalDateTime;
import java.util.List;
import java.util.Objects;
import java.util.concurrent.TimeUnit;
import java.util.function.Consumer;

/**
 * @author zhaoyifan
 * @date 2023/12/13 17:01
 */
@Service
@Slf4j
public class BaiduAiServiceImpl implements IBaiduAiService {

    private static final WebClient WEB_CLIENT = WebClient.builder()
            .defaultHeader(HttpHeaders.CONTENT_TYPE, org.springframework.http.MediaType.APPLICATION_JSON_VALUE)
            .build();

    public static final String API_KEY = "PwpOGSANB3AUPxQRDl2fP91Z";

    public static final String SECRET_KEY = "375YsGUe6nW6RGCsK5ORla4WRMp2gi35";

    static final OkHttpClient HTTP_CLIENT = new OkHttpClient().newBuilder()
            .connectTimeout(3600, TimeUnit.SECONDS) //连接超时
            .readTimeout(3600, TimeUnit.SECONDS) //读取超时
            .writeTimeout(3600, TimeUnit.SECONDS) //写超时
            .build();

    @Autowired
    private StringRedisTemplate redisTemplate;
    @Autowired
    IAiQuestionService questionService;
    @Autowired
    IAiQuestionItemService questionItemService;

    @Override
    public String callAi(List<AiMessage> aiMessageList) throws IOException {
        MediaType mediaType = MediaType.parse("application/json");
        JSONObject messages = new JSONObject();
        messages.put("messages",aiMessageList);
        System.out.println(JSONObject.toJSONString(messages));
        RequestBody body = RequestBody.create(mediaType, JSONObject.toJSONString(messages));
        Request request = new Request.Builder()
                .url("https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/completions_pro?access_token=" + getAccessToken())
                .method("POST", body)
                .addHeader("Content-Type", "application/json")
                .build();
        Response response = HTTP_CLIENT.newCall(request).execute();
        JSONObject result = JSONObject.parseObject(response.body().string());
        log.info("调用AI接口返回={}", result);
        return result.getString("result");
    }

    @Override
    public Flux<ChatResponse> callStream(List<AiMessage> aiMessageList) throws IOException {
        String completeUrl = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/completions_pro?access_token=" + getAccessToken();
        WebClient webClient = WEB_CLIENT.mutate().baseUrl(completeUrl).build();
        JSONObject messages = new JSONObject();
        messages.put("messages",aiMessageList);
        messages.put("stream",true);
        System.out.println(JSONObject.toJSONString(messages));

        return webClient.post()
                .body(BodyInserters.fromValue(messages))
                .accept(org.springframework.http.MediaType.TEXT_EVENT_STREAM)
                .retrieve()
                .bodyToFlux(ChatResponse.class)
                .doOnNext(this::handleErrResponse)
                .doOnError(WebClientResponseException.class, handleWebClientError());
    }

    @Override
    public AiMessageResponse callAiV2(AiMessageRequest messageRequest, Long userId) throws IOException {
        AiMessageResponse aiResponse = new AiMessageResponse();

        MediaType mediaType = MediaType.parse("application/json");
        JSONObject messages = new JSONObject();
        messages.put("messages",messageRequest.getList());
        System.out.println(JSONObject.toJSONString(messages));
        RequestBody body = RequestBody.create(mediaType, JSONObject.toJSONString(messages));
        Request request = new Request.Builder()
                .url("https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/completions_pro?access_token=" + getAccessToken())
                .method("POST", body)
                .addHeader("Content-Type", "application/json")
                .build();
        Response response = HTTP_CLIENT.newCall(request).execute();
        JSONObject result = JSONObject.parseObject(response.body().string());
        log.info("调用AI接口返回={}", result);
        if(Objects.nonNull(messageRequest.getQuestionId())){
            AiMessage aiMessage = messageRequest.getList().get(messageRequest.getList().size()-1);

            AiQuestionItem aiQuestionItem= new AiQuestionItem();
            aiQuestionItem.setContent(aiMessage.getContent());
            aiQuestionItem.setRole(aiMessage.getRole());
            aiQuestionItem.setQuestionId(messageRequest.getQuestionId());
            aiQuestionItem.setCreateDate(LocalDateTime.now());
            questionItemService.save(aiQuestionItem);

            AiQuestionItem aiQuestionItem1 = new AiQuestionItem();
            aiQuestionItem1.setContent(result.getString("result"));
            aiQuestionItem1.setRole("assistant");
            aiQuestionItem1.setQuestionId(messageRequest.getQuestionId());
            aiQuestionItem1.setCreateDate(LocalDateTime.now());
            questionItemService.save(aiQuestionItem1);

            aiResponse.setQuestionId(messageRequest.getQuestionId());
        }else{
            AiMessage aiMessage = messageRequest.getList().get(messageRequest.getList().size()-1);
            AiQuestion aiQuestion = new AiQuestion();
            aiQuestion.setTitle(aiMessage.getContent());
            aiQuestion.setUserId(userId);
            aiQuestion.setCreateDate(LocalDateTime.now());
            questionService.save(aiQuestion);

            AiQuestionItem aiQuestionItem= new AiQuestionItem();
            aiQuestionItem.setContent(aiMessage.getContent());
            aiQuestionItem.setRole(aiMessage.getRole());
            aiQuestionItem.setQuestionId(aiQuestion.getId());
            aiQuestionItem.setCreateDate(LocalDateTime.now());
            questionItemService.save(aiQuestionItem);

            AiQuestionItem aiQuestionItem1 = new AiQuestionItem();
            aiQuestionItem1.setContent(result.getString("result"));
            aiQuestionItem1.setRole("assistant");
            aiQuestionItem1.setQuestionId(aiQuestion.getId());
            aiQuestionItem1.setCreateDate(LocalDateTime.now());
            questionItemService.save(aiQuestionItem1);

            aiResponse.setQuestionId(aiQuestion.getId());
        }
        aiResponse.setContent(result.getString("result"));
        return aiResponse;
    }

    @Override
    public List<AiQuestion> chatList(Long userId) {
        List<AiQuestion> list = questionService.list(new LambdaQueryWrapper<AiQuestion>()
                .eq(AiQuestion::getUserId, userId)
        );
        if(Objects.nonNull(list)){
            list.forEach(aiQuestion -> {
                List<AiQuestionItem> itemList = questionItemService.list(new LambdaQueryWrapper<AiQuestionItem>()
                        .eq(AiQuestionItem::getQuestionId, aiQuestion.getId())
                );
                aiQuestion.setItemList(itemList);
            });
        }
        return list;
    }

    @Override
    public void updateChatTitle(AiQuestion aiQuestion) {
        aiQuestion.setUpdateDate(LocalDateTime.now());
        questionService.updateById(aiQuestion);
    }

    @Override
    public void deleteChat(AiQuestion aiQuestion) {
        questionService.removeById(aiQuestion);
    }

    /**
     * 从用户的AK，SK生成鉴权签名（Access Token）
     *
     * @return 鉴权签名（Access Token）
     * @throws IOException IO异常
     */
    public String getAccessToken() throws IOException {
        String key = "BAIDU_AI_TOKEN";
        if(redisTemplate.hasKey("BAIDU_AI_TOKEN")){
            return redisTemplate.opsForValue().get(key);
        }else{
            MediaType mediaType = MediaType.parse("application/x-www-form-urlencoded");
            RequestBody body = RequestBody.create(mediaType, "grant_type=client_credentials&client_id=" + API_KEY
                    + "&client_secret=" + SECRET_KEY);
            Request request = new Request.Builder()
                    .url("https://aip.baidubce.com/oauth/2.0/token")
                    .method("POST", body)
                    .addHeader("Content-Type", "application/x-www-form-urlencoded")
                    .build();
            Response response = HTTP_CLIENT.newCall(request).execute();
            JSONObject jsonObject = JSONObject.parseObject(response.body().string());
            log.info("获取token返回={}", jsonObject);
            redisTemplate.opsForValue().set(key, jsonObject.getString("access_token"), jsonObject.getLong("expires_in") , TimeUnit.SECONDS);
            return jsonObject.getString("access_token");
        }
    }

    private <T> void handleErrResponse(T response) {
        if (response instanceof ChatResponse) {
            ChatResponse chatResponse = (ChatResponse) response;
            log.info("流式请求返回={}",JSONObject.toJSONString(chatResponse));
            if (chatResponse.getErrorCode() != null) {
                ErrorResponse errorResponse = ErrorResponse.builder()
                        .id(chatResponse.getId())
                        .logId(chatResponse.getLogId())
                        .ebCode(chatResponse.getEbCode())
                        .errorMsg(chatResponse.getErrorMsg())
                        .errorCode(chatResponse.getErrorCode())
                        .build();
                throw new RuntimeException(errorResponse.getErrorMsg());
            }
        }
    }

    private static Consumer<Throwable> handleWebClientError() {
        return err -> {
            log.error("请求错误 => {} {}", err instanceof WebClientResponseException
                    ? ((WebClientResponseException) err).getStatusCode()
                    : "Unknown", err.getMessage());
            throw new RuntimeException("系统网络异常");
        };
    }
}
